package milvus;

import cn.hutool.core.collection.CollectionUtil;
import com.google.gson.Gson;
import com.google.gson.JsonObject;
import io.milvus.v2.client.ConnectConfig;
import io.milvus.v2.client.MilvusClientV2;
import io.milvus.v2.common.DataType;
import io.milvus.v2.common.IndexParam;
import io.milvus.v2.service.collection.request.AddFieldReq;
import io.milvus.v2.service.collection.request.CreateCollectionReq;
import io.milvus.v2.service.collection.request.DropCollectionReq;
import io.milvus.v2.service.collection.request.GetLoadStateReq;
import io.milvus.v2.service.collection.response.ListCollectionsResp;
import io.milvus.v2.service.database.request.CreateDatabaseReq;
import io.milvus.v2.service.database.request.DropDatabaseReq;
import io.milvus.v2.service.database.response.ListDatabasesResp;
import io.milvus.v2.service.vector.request.GetReq;
import io.milvus.v2.service.vector.request.InsertReq;
import io.milvus.v2.service.vector.request.SearchReq;
import io.milvus.v2.service.vector.request.data.FloatVec;
import io.milvus.v2.service.vector.response.GetResp;
import io.milvus.v2.service.vector.response.InsertResp;
import io.milvus.v2.service.vector.response.QueryResp;
import io.milvus.v2.service.vector.response.SearchResp;
import lombok.Data;

import java.io.Serializable;
import java.util.ArrayList;
import java.util.Arrays;
import java.util.Collections;
import java.util.List;

@Data
public class milvus implements Serializable {


    // 1. Connect to Milvus server
    private static ConnectConfig connectConfig = ConnectConfig.builder()
            .uri("http://116.62.142.81:19530")
            .token("root:Milvus")
            .build();
    private static MilvusClientV2 client = new MilvusClientV2(connectConfig);
    public static void main(String[] args) throws InterruptedException {

        createDataBase("cheshi");
        createCollection("cheshi_1");
//        deleteCollection("cheshi_1");
//        deleteDatabase("cheshi");
//        addData("cheshi_1");
        query2("cheshi_1");
    }



    private static void createDataBase(String databaseName) throws InterruptedException {
        ListDatabasesResp listDatabasesResp = client.listDatabases();

        if (CollectionUtil.isNotEmpty(listDatabasesResp.getDatabaseNames())){
            if (listDatabasesResp.getDatabaseNames().contains(databaseName)){
                client.useDatabase(databaseName);
                return;
            }
        }

        // 创建数据库
        CreateDatabaseReq createDatabaseParam = CreateDatabaseReq.builder()
                .databaseName(databaseName)
                .build();
        client.createDatabase(createDatabaseParam);
        client.useDatabase(databaseName);

    }

    private static void deleteDatabase(String databaseName) {
        DropDatabaseReq dropDatabaseReq = DropDatabaseReq.builder()
                .databaseName(databaseName)
                .build();
        client.dropDatabase(dropDatabaseReq);
    }
    private static void createCollection( String collectionName) {

        ListCollectionsResp listCollectionsResp = client.listCollections();
        if (CollectionUtil.isNotEmpty(listCollectionsResp.getCollectionNames()) &&listCollectionsResp.getCollectionNames().contains(collectionName)){
            return;
        }
        //创建模式
        CreateCollectionReq.CollectionSchema schema = client.createSchema();
        // 3.2 添加字段
        schema.addField(AddFieldReq.builder()
                .fieldName("id")
                .dataType(DataType.Int64)
                .isPrimaryKey(true)
                .autoID(false)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("vector")
                .dataType(DataType.FloatVector)
                .dimension(5)
                .build());

        schema.addField(AddFieldReq.builder()
                .fieldName("color")
                .dataType(DataType.VarChar)
                .maxLength(512)
                .build());



        // 创建索引
        IndexParam indexParamForIdField = IndexParam.builder()
                .fieldName("id")
                .indexType(IndexParam.IndexType.STL_SORT)
                .build();

        IndexParam indexParamForVectorField = IndexParam.builder()
                .fieldName("vector")
                .indexType(IndexParam.IndexType.AUTOINDEX)
                .metricType(IndexParam.MetricType.COSINE)
                .build();

        List<IndexParam> indexParams = new ArrayList<>();
        indexParams.add(indexParamForIdField);
        indexParams.add(indexParamForVectorField);


        // 3.4 创建具有模式和索引参数的集合
        CreateCollectionReq customizedSetupReq1 = CreateCollectionReq.builder()
                .collectionName(collectionName)
                .collectionSchema(schema)
                .indexParams(indexParams)
                .build();

        client.createCollection(customizedSetupReq1);

        // 3.5 获取集合的加载状态
        GetLoadStateReq customSetupLoadStateReq1 = GetLoadStateReq.builder()
                .collectionName(collectionName)
                .build();

        Boolean loaded = client.getLoadState(customSetupLoadStateReq1);
        System.out.println(loaded);

    }
    private static void deleteCollection(String collectionName) {
        DropCollectionReq collectionNameReq = DropCollectionReq.builder()
                .collectionName(collectionName)
                .build();
        client.dropCollection(collectionNameReq);
    }
    private static void addData(String collectionName) {
        Gson gson = new Gson();
        List<JsonObject> data = Arrays.asList(
                gson.fromJson("{\"id\": 0, \"vector\": [0.3580376395471989f, -0.6023495712049978f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f], \"color\": \"pink_8682\"}", JsonObject.class),
                gson.fromJson("{\"id\": 1, \"vector\": [0.19886812562848388f, 0.06023560599112088f, 0.6976963061752597f, 0.2614474506242501f, 0.838729485096104f], \"color\": \"red_7025\"}", JsonObject.class),
                gson.fromJson("{\"id\": 2, \"vector\": [0.43742130801983836f, -0.5597502546264526f, 0.6457887650909682f, 0.7894058910881185f, 0.20785793220625592f], \"color\": \"orange_6781\"}", JsonObject.class),
                gson.fromJson("{\"id\": 3, \"vector\": [0.3172005263489739f, 0.9719044792798428f, -0.36981146090600725f, -0.4860894583077995f, 0.95791889146345f], \"color\": \"pink_9298\"}", JsonObject.class),
                gson.fromJson("{\"id\": 4, \"vector\": [0.4452349528804562f, -0.8757026943054742f, 0.8220779437047674f, 0.46406290649483184f, 0.30337481143159106f], \"color\": \"red_4794\"}", JsonObject.class),
                gson.fromJson("{\"id\": 5, \"vector\": [0.985825131989184f, -0.8144651566660419f, 0.6299267002202009f, 0.1206906911183383f, -0.1446277761879955f], \"color\": \"yellow_4222\"}", JsonObject.class),
                gson.fromJson("{\"id\": 6, \"vector\": [0.8371977790571115f, -0.015764369584852833f, -0.31062937026679327f, -0.562666951622192f, -0.8984947637863987f], \"color\": \"red_9392\"}", JsonObject.class),
                gson.fromJson("{\"id\": 7, \"vector\": [-0.33445148015177995f, -0.2567135004164067f, 0.8987539745369246f, 0.9402995886420709f, 0.5378064918413052f], \"color\": \"grey_8510\"}", JsonObject.class),
                gson.fromJson("{\"id\": 8, \"vector\": [0.39524717779832685f, 0.4000257286739164f, -0.5890507376891594f, -0.8650502298996872f, -0.6140360785406336f], \"color\": \"white_9381\"}", JsonObject.class),
                gson.fromJson("{\"id\": 9, \"vector\": [0.5718280481994695f, 0.24070317428066512f, -0.3737913482606834f, -0.06726932177492717f, -0.6980531615588608f], \"color\": \"purple_4976\"}", JsonObject.class)
        );

        InsertReq insertReq = InsertReq.builder()
                .collectionName(collectionName)
                .data(data)
                .build();
        InsertResp insertResp = client.insert(insertReq);
        System.out.println(insertResp);

    }

    private static void query(String collectionName) {
        GetReq getReq = GetReq.builder()
                .collectionName(collectionName)
                .ids(Arrays.asList(0, 1, 2))
                .outputFields(Arrays.asList("vector", "color"))
                .build();

        GetResp getResp = client.get(getReq);

        List<QueryResp.QueryResult> results = getResp.getGetResults();
        for (QueryResp.QueryResult result : results) {
            System.out.println(result.getEntity());
        }
    }

    private static void query2(String collectionName) {
        FloatVec queryVector = new FloatVec(new float[]{0.3580376395471989f, 0.06023560599112088f, 0.18414012509913835f, -0.26286205330961354f, 0.9029438446296592f});
        SearchReq searchReq = SearchReq.builder()
                .collectionName(collectionName)
                .data(Collections.singletonList(queryVector))
                .topK(10)
                .build();

        SearchResp searchResp = client.search(searchReq);

        List<List<SearchResp.SearchResult>> searchResults = searchResp.getSearchResults();
        for (List<SearchResp.SearchResult> results : searchResults) {
            System.out.println("TopK results:");
            for (SearchResp.SearchResult result : results) {
                System.out.println(result);
            }
        }
    }
}
